package com.inderlab.jpin;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.io.InputStream;
import java.lang.Runtime;

import javassist.NotFoundException;
import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtField;
import javassist.Modifier;
import javassist.CtMethod;
import javassist.CtNewMethod;
import javassist.expr.ExprEditor;

import com.inderlab.jpin.JPEditor;
import com.inderlab.jpin.JPColor;
import com.inderlab.jpin.JPConfig;
import com.inderlab.jpin.JPCodeBook;

public class JPFormer implements ClassFileTransformer
{
	private static boolean compactmode = com.inderlab.jpin.JPAgent.compactmode;
	private static boolean verbosemode = com.inderlab.jpin.JPAgent.verbosemode;
	private static boolean sidemode    = com.inderlab.jpin.JPAgent.sidemode;
	private static boolean matchmode   = com.inderlab.jpin.JPAgent.matchmode;
	private static boolean depthmode   = com.inderlab.jpin.JPAgent.depthmode;
	
	private ExprEditor editor = new JPEditor();
	private JPConfig config = null;
	
	public void setConfig(JPConfig poleconfig)
	{
		config = poleconfig;
	}
	
	public byte[] transform(ClassLoader loader, 
							String className,
							Class<?> classBeingRedefined, 
							ProtectionDomain protectionDomain,
							byte[] classfileBuffer)
							throws IllegalClassFormatException
	{
		byte[] transformed = classfileBuffer;
					
		if(verbosemode)
			if(matchmode)
			System.out.println((sidemode ? JPColor.MARGIN:"")+JPColor.C4+"[jpin]  [Class] " + loader + " ["+className+"] size="+classfileBuffer.length+JPColor.RESET);
			
		return transform(className, transformed);
	}
	
	//
	// Transform all methods for a class
	//
	protected byte[] transform(String className, byte[] byteCode)
	{
		InputStream is = null;
		byte[] modifiedBytecode = byteCode;
		
		// Before exiting JVM, jpin saves all info.
		/*
		if(className.equals("java/lang/Shutdown"))
		{
			// By Shutdown hook due to the GC before a Shutdown method works.
			// JPCodeBook.writeCodeBook();		
		}
		*/
		
		// Find and Transform a filtered method
		try {
			
			if(config.hasClass(className))
			{
				ClassPool classPool = ClassPool.getDefault();
				is = new ByteArrayInputStream(byteCode);
				CtClass cc = classPool.makeClass(is);
				
				// Ignore Interface
				if(cc.isInterface()) return byteCode;

				//-----------------------------------------
				// Insert Thread Local Variable to each Class (Only one for each class)
				//-----------------------------------------
				
				try {
					cc.getField("threadLocal_pin");
				} 
				catch(NotFoundException e)
				{
					/* --Due to polestar error: modifying fields
					CtField thlocal = new CtField(classPool.get("java.lang.ThreadLocal"),"threadLocal_pin", cc);
					thlocal.setModifiers(Modifier.PUBLIC | Modifier.STATIC);
					cc.addField(thlocal,"new ThreadLocal()");
					*/
					
					//
					// How about to make Hashtable for methods of a class to maintain the number of callings and memory profile
					//
					
					/*	
					cc.addField(thlocal,"null");
					CtMethod m = CtNewMethod.make("public void pin_tls()"+
						"{"+
						" if(threadLocal_pin == null)"+
						" 		threadLocal_pin = new ThreadLocal();"+
						" threadLocal_pin.set( new String(\"Test\") );"+
						" System.out.println(threadLocal_pin.get());"+
						"}"
						,cc);
					cc.addMethod(m);							
					*/
				}
				
				//-----------------------------------------
				// Put Timing Code to All Methods 
				//-----------------------------------------
				if(verbosemode)
					if(matchmode)
					System.out.println((sidemode ? JPColor.MARGIN:"")+"[");
				
				CtMethod[] methods = cc.getDeclaredMethods();
				for(int i=0; i<methods.length; i++)
				{						
					// Specific
					if(!methods[i].getName().equals("pin_tls"))
					if(config.hasMethod(className, methods[i].getName()))
					{
						if(verbosemode)
							if(matchmode)
							System.out.println((sidemode ? JPColor.MARGIN:"")+"    * "+className+"."+methods[i].getName());

						// Add to Codebook:
						// Assign code# to className.methodName
						JPCodeBook.register(className, methods[i].getName());
																				
						// Insert Time and Memory Code to each Method
						addMemTime(cc, methods[i].getName());						
					}
				}

				modifiedBytecode = cc.toBytecode();
				if(verbosemode)
					if(matchmode)
					System.out.println((sidemode ? JPColor.MARGIN:"")+"] size="+modifiedBytecode.length+JPColor.RESET);												
			}
			
			return modifiedBytecode;
			
		} catch(RuntimeException e) {
			throw e;
		} catch(Exception e) {
			throw new RuntimeException(e);
		} finally {
			if(is != null) try { is.close(); } catch(Exception e) 
			{
				if(verbosemode)
					System.out.println((sidemode ? JPColor.MARGIN:"")+"[jpin] [ is.close() Exception ]--");
			}
		}
		
	}

	//
	// Add Time Measuring Code to a Method
	//
	private void addMemTime(CtClass clas, String mname)
	        throws NotFoundException, CannotCompileException 
	{			
        CtMethod mold = clas.getDeclaredMethod(mname);
        String nname = mname+"$impl";
        mold.setName(nname);
        CtMethod mnew = CtNewMethod.copy(mold, mname, clas, null);

        String type = mold.getReturnType().getName();
        StringBuffer body = new StringBuffer();
        body.append("{ long pin_sttime = System.nanoTime(); ");
		body.append("int callid = com.inderlab.jpin.JPCodeBook.callcounter++;");
		body.append("Runtime rt=Runtime.getRuntime(); long hfs = rt.freeMemory(); ");

        // Access Violation
		// body.append("System.out.println(\"threadlocal=\"+Thread.currentThread().threadLocals);");			
	
		//-----------------------------------------
		// Trace Method Invocation Stack's Frames
		//-----------------------------------------
		//   Method 1. From Throwable (Exception)
		//   Method 2. Thread's Stack Trace (Thread.currentThread().getStackTrace())
		//
		//   * body.append("StackTraceElement[] stackElements = new Throwable().getStackTrace();\n");
		//   * Thread.getAllStackTraces()
		//
		// Need to be analyzed the overhead of getting Strack Trace
		//-----------------------------------------
					
		body.append("StackTraceElement[] stackElements = Thread.currentThread().getStackTrace();");
		body.append("int seidx = stackElements.length >= 5 ? 4:2;");
		
		// TODO: hashCode() to HEX Code
		// body.append("String ori_caller=stackElements[seidx].getClassName()+\".\"+stackElements[seidx].getMethodName()+\" \"+(stackElements[seidx].hashCode()&0x7fffffff);");
		body.append("String ori_caller = stackElements[seidx].getClassName()+\".\"+stackElements[seidx].getMethodName();");
		
		if(verbosemode)
		{
			if(compactmode) {
				// Search Codebook: Caller and Callee
	 			body.append("int cc1 = com.inderlab.jpin.JPCodeBook.getCode(ori_caller);");
	 			body.append("int cc2 = com.inderlab.jpin.JPCodeBook.getCode(\""+clas.getName()+"."+mname+"\");");
				body.append("System.out.print(\""+(sidemode ? JPColor.MARGIN:"")+JPColor.C2+"\");");
				body.append("if(com.inderlab.jpin.JPAgent.depthmode) for(int i=0; i<stackElements.length-3; i++) System.out.print(\" \");");
				body.append("System.out.println(\""+JPColor.C2+"[>] \"+callid+\" \"+cc1+\" \"+cc2+\" \"+stackElements.length+\" \"+((pin_sttime/1000000)&0x0000ffff)+\""+JPColor.RESET+"\");");				
//				body.append("System.out.println(\""+JPColor.C2+"[>] \"+callid+\" \"+cc1+\" \"+cc2+\" \"+stackElements.length+\" \"+((pin_sttime>>19)&0x0000ffff)+\""+JPColor.RESET+"\");");				
			}
			else {
				// TODO: Compact Code
				body.append("System.out.print(\""+(sidemode ? JPColor.MARGIN:"")+JPColor.C2+"\");");
				body.append("if(com.inderlab.jpin.JPAgent.depthmode) for(int i=0; i<stackElements.length-3; i++) System.out.print(\" \");");
				body.append("System.out.println(\"[>] \"+callid+\" \"+ori_caller+\" -> "+clas.getName()+"."+mname+" \"+stackElements.length+\" \"+((pin_sttime/1000000)&0x0000ffff)+\""+JPColor.RESET+"\");");
//				body.append("System.out.println(\"[>] \"+callid+\" \"+ori_caller+\" -> "+clas.getName()+"."+mname+" \"+stackElements.length+\" \"+((pin_sttime>>19)&0x0000ffff)+\""+JPColor.RESET+"\");");

			}			
		}		


		if(mname.equals("registerMBean")) {
			body.append("System.out.println(\"  [>] obj=\"+$1+\"\\n  [>] name=\"+$2);");			
		} else
		if(mname.equals("invoke")) {
			body.append("System.out.println(\"  [>] name=\"+$1+\"\\n  [>] method=\"+$2);");			
		} else
		if(mname.equals("queryMBeans") || mname.equals("queryNames")) {
			body.append("System.out.println(\"  [>] name=\"+$1);");			
		}
		
		//body.append("System.out.println(\""+JPColor.MARGIN+JPColor.C2+"  [jpin*] \"+ori_caller+\" -> "+clas.getName()+"."+mname+" \"+stackElements.length+\""+JPColor.RESET+"\");");
					
        if (!"void".equals(type)) 
		{
			body.append(type + " result = ");
		}

        body.append(nname + "($$);\n");

		//-----------------------------------------
		// Extra Memory Info
		//-----------------------------------------
		
        /* 
		body.append("long curr_mem_fs  = rt.freeMemory()-hfs;\n");
		body.append("long curr_mem_tot = rt.totalMemory();\n");
		body.append("long curr_mem_max = rt.maxMemory();\n");
		body.append("System.out.println(\""+JPColor.MARGIN+JPColor.C5+"  [jpin] <- "+clas.getName()+"."+mname+
			" [\" + ((System.nanoTime()-start)>>19) + " + "\",\"+curr_mem_fs+\",\"+curr_mem_tot+\",\"+curr_mem_max	+\"]"+JPColor.RESET+"\");\n");
		*/
		
		if(verbosemode)
		{
			if(compactmode) {
				body.append("System.out.print(\""+(sidemode ? JPColor.MARGIN:"")+JPColor.C5+"\");");
				body.append("if(com.inderlab.jpin.JPAgent.depthmode) for(int i=0; i<stackElements.length-3; i++) System.out.print(\" \");");
		        body.append("System.out.println(\"[<] \"+callid+\" \"+cc1+\" \"+cc2+\" \"+((pin_sttime/1000000)&0x0000ffff)+ \" \" + ((System.nanoTime()-pin_sttime)/1000000) + " + "\" \"+(rt.freeMemory()-hfs)+\" "+JPColor.RESET+"\");");
//		        body.append("System.out.println(\"[<] \"+callid+\" \"+cc1+\" \"+cc2+\" \"+((pin_sttime>>19)&0x0000ffff)+ \" \" + ((System.nanoTime()-pin_sttime)>>19) + " + "\" \"+(rt.freeMemory()-hfs)+\" "+JPColor.RESET+"\");");
			} else			
			{
				// TODO: Compact Code
				body.append("System.out.print(\""+(sidemode ? JPColor.MARGIN:"")+JPColor.C5+"\");");
				body.append("if(com.inderlab.jpin.JPAgent.depthmode) for(int i=0; i<stackElements.length-3; i++) System.out.print(\" \");");
	        	body.append("System.out.println(\"[<] \"+callid+\" \"+ori_caller+\" <- "+clas.getName()+"."+mname+
   					" \" + ((pin_sttime/1000000)&0x0000ffff)+\" \"+((System.nanoTime()-pin_sttime)/1000000) + " + "\" \"+(rt.freeMemory()-hfs)+\""+JPColor.RESET+"\");");				
//   				" \" + ((pin_sttime>>19)&0x0000ffff)+\" \"+((System.nanoTime()-pin_sttime)>>19) + " + "\" \"+(rt.freeMemory()-hfs)+\""+JPColor.RESET+"\");");				

			}
		}
		
		// body.append("System.out.println(\""+JPColor.MARGIN+JPColor.C5+"  [jpin*] \"+ori_caller+\" <- "+clas.getName()+"."+mname+
		//	   			" [\" + ((System.nanoTime()-start)>>19) + " + "\",\"+(rt.freeMemory()-hfs)+\"]"+JPColor.RESET+"\");");

		//-----------------------------------------
		// Store to Thread Local Storage
		//-----------------------------------------

		//body.append("threadLocal_pin.set(new String(\"test\"));");

		//-----------------------------------------
		// Return
		//-----------------------------------------

        if (!"void".equals(type)) 
            body.append("return result;\n");

        body.append("}");
        mnew.setBody(body.toString());
        clas.addMethod(mnew);

        //print the generated code block just to show what was done
        //System.out.println("Interceptor method body:");
        //System.out.println(body.toString());
	}	
}
